import math
import numpy as np
from numpy import random
import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.quantization import fake_quant_with_min_max_vars

n, ic, h, w = 1, 128, 480, 480
oc, _, kh, kw = 128, 128, 3, 3

# n, ic, h, w = 1, 96, 16, 256
# oc, _, kh, kw = 80, 96, 3, 3
input = tf.placeholder(name="input", dtype=tf.float32, shape=(n, h, w, ic))    # NHWC
weight0 = tf.get_variable("weight0", dtype=tf.float32, shape=(kh, kw, ic, oc))   # HWIO
bias0 = tf.get_variable("bias0", dtype=tf.float32, shape=(oc,))
conv0 = tf.nn.conv2d(input, weight0, strides=1, padding='SAME', use_cudnn_on_gpu=False, name="conv0")
output = tf.nn.bias_add(conv0, bias0, name="bias0")
output = tf.identity(output, name="output")

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # feed_dict = {input}
    # weight_np = sess.run(weight0)
    # bias_np = sess.run(bias0)
    

    # export pb: TensorFlow保存模型为PB文件 https://zhuanlan.zhihu.com/p/32887066
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['output'])
    with tf.gfile.FastGFile('/home/cqchu/models/conv_fp32.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

    # export float32 tflite
    converter = tf.lite.TFLiteConverter.from_session(sess, [input], [output])
    open("/home/cqchu/models/conv_fp32.tflite", "wb").write(converter.convert())

    # export quantized tflite
    # converter = tf.lite.TFLiteConverter.from_session(sess, [input], [output])
    # converter.inference_type = tf.lite.constants.QUANTIZED_UINT8
    # converter.quantized_input_stats = {"input": (128, 128)}
    # converter.default_ranges_stats = (0, 1)
    # open("uint8.tflite", "wb").write(converter.convert())

# for tv in tf.trainable_variables():
#     print(tv)
# bias = tf.get_default_graph().get_tensor_by_name("bias0:0")
# print(bias)

    